import catboost as cb
from sklearn.metrics import mean_absolute_error,  make_scorer
import pandas as pd
import numpy as np
import nni

def main(args):
    cat = {
        'model': 'category',
        'brand': 'category',
        'bodyType': 'category',
        'fuelType': 'category',
        'gearbox': 'category',
        'notRepairedDamage': 'category',
    }

    data0 = pd.read_csv('./user_data/df_s.csv', sep=' ', dtype=cat).drop(
        ['SaleID', 'regionCode'], axis=1)
    train_X = data0[data0.train == 1].drop(['price'], axis=1)
    train_y = data0[data0.train == 1]['price']
    train_y_ln = np.log1p(train_y)

    from sklearn.model_selection import cross_val_score
    n = 1

    def maee(y_true, y_pred):
        nonlocal n
        loss = mean_absolute_error(np.expm1(y_true), np.expm1(y_pred))
        print(f'mae in fold {n}:{loss}')
        n += 1
        return loss
    model = cb.CatBoostRegressor(**args)
    fit_params = {
        'cat_features': [i for i in cat.keys() if i in train_X.columns],
        'early_stopping_rounds': 300,
        'verbose': 300
    }
    mae = cross_val_score(model, X=train_X, y=train_y_ln, verbose=0,
                          cv=5, scoring=make_scorer(maee), fit_params=fit_params)
    print(mae)
    print(f'with mean mae:{np.mean(mae)}')


if __name__ == '__main__':
    params = {
        #'n_estimators': 1000000,
        'loss_function': 'MAE',
        'eval_metric': 'MAE',
        # 'learning_rate': 0.02,# 设置为0.02会导致训练时间极——————————————长，建议设高一点或者找gpu加速
        'depth': 6,
        #'use_best_model': True,
        'subsample': 0.6,
        'bootstrap_type': 'Bernoulli',
        'reg_lambda': 3,
        'one_hot_max_size': 2,
    }
    main(params)
